# -*- coding: utf-8 -*-
from __future__ import annotations
import json
import sys
from typing import Tuple

from llm_provider import chat
from tools import dispatch_tool
from prompts import AGENT_SYSTEM, FORMAT_EXAMPLE

def _strip_fence(s: str) -> str:
    # 去掉 ```json ... ``` 之类代码块，容忍模型偶发包裹
    s = s.strip()
    if s.startswith("```"):
        s = s.strip("`")
        # 简单处理：截取第一个换行后的部分
        parts = s.split("\n", 1)
        if len(parts) == 2:
            s = parts[1]
    return s.strip()

def _parse_json(s: str) -> dict:
    s = _strip_fence(s)
    return json.loads(s)

def run_agent(user_input: str, max_steps: int = 3) -> str:
    """最小单工具 Agent 主循环。返回 final 文本。"""
    messages = [
        {"role": "system", "content": AGENT_SYSTEM},
        {"role": "user", "content": FORMAT_EXAMPLE},
        {"role": "user", "content": user_input},
    ]

    for _ in range(max_steps):
        out = chat(messages, json_mode=True)
        try:
            data = _parse_json(out)
        except json.JSONDecodeError:
            # 强制提醒：只输出 JSON
            messages.append({"role": "assistant", "content": "输出格式错误，请只输出 JSON 对象。"})
            continue

        if "tool" in data:
            name = data["tool"]
            args = data.get("args", {}) or {}
            try:
                print(f"[Debug: 执行工具：{name}，参数：{args}]")
                result = dispatch_tool(name, args)
            except Exception as e:
                # 把错误作为“观察”反馈给模型，提示修正
                messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
                messages.append({"role": "user", "content": f"OBS: 工具执行失败，原因：{e}. 请修正参数或改用 final。"})
                continue

            # 把工具结果作为观察提供，提示收敛为 final
            messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
            messages.append({"role": "user", "content": f'OBS: 工具 {name} 返回 {result}. 若可给出答案，请输出 {{"final": "..."}}'})
            continue

        if "final" in data:
            return str(data["final"])

        # 若既无 tool 又无法终止，要求重试
        messages.append({"role": "assistant", "content": "缺少 tool 或 final 字段，请按约定输出。"})
    raise RuntimeError("达到最大步数仍未获得 final。")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print('用法: python agent.py "2+3是多少？"')
        sys.exit(1)
    query = " ".join(sys.argv[1:]).strip()
    try:
        ans = run_agent(query)
        print(ans)
    except Exception as e:
        print(f"[Error] {e}")
        sys.exit(2)
